package org.simtube.feedback;

import java.util.ArrayList;

import org.simtube.util.*;

/*
 * 
 */
public class Feedback {
	
	private DocItem[] dvectors;
	private int dims;
	
	public Feedback(){
	}
	/**
	 * constructor
	 * @param data  document vectors
	 */
	public Feedback(DocItem[] data)
	{
		this.dvectors = data;
		this.dims = data[0].vector.getDim();
	}
	
	public void setDvectors(DocItem[] dvectors) {
		this.dvectors = dvectors;
		this.dims = dvectors[0].vector.getDim();
	}

	/**
	 * constructor
	 * @param data  document vectors
	 * @param _dims   dimension of the vectors
	 */
	public Feedback(DocItem[] data, int _dims)
	{
		this(data);
		if (_dims > 0) this.dims = _dims;
		else throw new IllegalArgumentException("dims <= 0");
	}

	/**
	 * Get a new query vector according to the idea of relevant feedback
	 * using Ide-DecHi.
	 * @param q   original query vector
	 * @param dr  relevant doc id array. can be null
	 * @param dn  non-relevant doc id array. can be null
	 * @return   a new query vector, null = exception
	 */
	public DocVector getQuery(DocVector q, int[] dr, int[] dn)
	{
		try
		{
			int i;
			DocVector result, sum;
			result = q.multiply(1);
			if (dr != null && dr.length > 0)
			{
				sum = result.multiply(0);
				for (i = 0; i < dr.length; i++)
					sum = sum.add(this.dvectors[dr[i]].vector);
				sum = sum.multiply(0.75);
				result = result.add(sum);
				
			}
			if (dn != null && dn.length > 0)
			{
				double min = q.norm();
				double temp;
				sum = null;
				for (i = 0; i < dn.length; i++)
				{
					temp = q.multiply(this.dvectors[dn[i]].vector);
					temp = temp / this.dvectors[dn[i]].vector.norm();
					if (temp < min)
					{
						min = temp;
						sum = this.dvectors[dn[i]].vector;
					}
				}
				if (min < q.norm())
				{
					sum = sum.multiply(0.25);
					result = result.minus(sum);
				}
			}
			return result.normalize();
		}
		catch (Exception e)
		{
			e.printStackTrace();
			return null;
		}		
	}
	
	/**
	 * Use relevance feedback to reweight terms
	 * @param query  query vector
	 * @param dr    vectors of relevant docs, can be null
	 * @param dn    vectoers of non-relevant docs, can be null
	 * @return    list of terms that should be reweighted, the tf have no meaning
	 */
	public ArrayList<TermTFIDF> reweightTerms(ArrayList<TermTFIDF> query,
			ArrayList<TermTFIDF>[] dr, ArrayList<TermTFIDF>[] dn, boolean expand)
	{
		try
		{
			if ((dr == null || dr.length == 0) && (dn == null || dn.length == 0))
				return new ArrayList<TermTFIDF>();
			class Term{
				int index;
				int rcount;
				int ncount;
				
				public Term(int _index)
				{
					rcount = ncount = 0;
					this.index = _index;
				}
				
				Term isIn(ArrayList<Term> l, TermTFIDF t)
				{
					for (Term e : l)
						if (e.index == t.getTermIndex())
							return e;
					return null;
				}
			}
			ArrayList<Term> data = new ArrayList<Term>();
			data.clear();
			Term temp;
			Term func = new Term(-1); // just use for function
			
			if (dr != null)
			{
				for (ArrayList<TermTFIDF> d : dr)
				{
					for (TermTFIDF t : d)
					{
						temp = func.isIn(data, t);
						if (temp != null) temp.rcount++;
						else data.add(new Term(t.getTermIndex()));
					}
				}
			}
			
			if (dn != null)
			{
				for (ArrayList<TermTFIDF> d : dn)
				{
					for (TermTFIDF t : d)
					{
						temp = func.isIn(data, t);
						if (temp != null) temp.ncount++;
						else data.add(new Term(t.getTermIndex()));
					}
				}
			}
			
			ArrayList<TermTFIDF> result = new ArrayList<TermTFIDF>();
			result.clear();
			double newIDF, p, q, emax, expRank;
			TermTFIDF newTerm, expandTerm = null;
			emax = 0;
			// the relevant count threshold of the expansion term
			
			for (Term t : data)
			{
				// calculate new IDF
				p = t.rcount + 0.5;
				p /= (dr == null? 0 : dr.length) + 1;
				q = t.ncount + 0.5;
				q /= (dn == null? 0 : dn.length) + 1;
				newIDF = Math.log(p * (1-q) / q / (1 - p));
				
				// add to modify list
				newTerm = new TermTFIDF(t.index, 0, newIDF);
				result.add(newTerm);
				
				// expansion
				if (expand)
				{
					expRank = newIDF * (p - q);
					if (expRank > emax)
					{
						expandTerm = newTerm;
						emax = expRank;
					}
				}
			}
			
			if (expandTerm != null)	query.add(expandTerm);
			return result;
		}
		catch (Exception e)
		{
			e.printStackTrace();
			return null;
		}
	}
	
	/*
	 * GETTERs and SETTERs below
	 */
	
	public int getDims()
	{
		return this.dims;
	}
	
	public int setDims(int _dims)
	{
		if (_dims > 0) this.dims = _dims;
		return this.dims;
	}
}
